import numpy as np

np.random.seed(1337)
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D
import MINIST

batch_size = 128
nb_classes = 10
epochs = 10
img_rows, img_cols = 28, 28
nb_filters = 32
pool_size = (2, 2)
kernel_size = (5, 5)
input_shape = (img_rows, img_cols, 1)
X_train, Y_train = MINIST.get_training_data_set(6000, False)
X_test, Y_test = MINIST.get_test_data_set(1000, False)
X_train = np.array(X_train).astype(bool).astype(float) / 255
X_train = X_train[:, :, :, np.newaxis]
Y_train = np.array(Y_train)

X_test = np.array(X_test).astype(bool).astype(float) / 255
X_test = X_test[:, :, :, np.newaxis]
Y_test = np.array(Y_test)
print('样本数据集维度：', X_train.shape, Y_train.shape)
print('测试数据集维度：', X_test.shape, Y_test.shape)
model = Sequential()
model.add(Conv2D(6, kernel_size, input_shape=input_shape, strides=1))

model.add(AveragePooling2D(pool_size=pool_size, strides=2))
model.add(Conv2D(12, kernel_size, strides=1))
model.add(AveragePooling2D(pool_size=pool_size, strides=2))
model.add(Flatten())
model.add(Dense(nb_classes))
model.add(Activation('sigmoid'))

model.compile(loss='categorical_crossentropy', optimizer='adadelta',metrics = ['accuracy'])

model.fit(X_train, Y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_data=(X_test, Y_test))

score = model.evaluate(X_test, Y_test, verbose=0)
print('TEST SCORE:', score[0])
print('test accuracy:', score[1])
model.save('cnn_model.h5')
